from mcp.server.fastmcp import FastMCP

# 创建MCP服务实例
mcp = FastMCP("Calculator Service")

@mcp.tool()
def add(a: float, b: float) -> str:
    """计算两数之和"""
    return str(a + b)

@mcp.tool()
def subtract(a: float, b: float) -> str:
    """计算两数之差"""
    return str(a - b)

@mcp.tool()
def multiply(a: float, b: float) -> str:
    """计算两数之积"""
    from decimal import Decimal, getcontext
    getcontext().prec = 50  # 设置足够高的精度
    return str(Decimal(str(a)) * Decimal(str(b)))

@mcp.tool()
def divide(a: float, b: float) -> str:
    """计算两数之商"""
    try:
        if b == 0:
            return "错误: 除数不能为零"
        return str(a / b)
    except Exception as e:
        return f"计算错误: {str(e)}"

@mcp.tool()
def nth_prime(k: int) -> str:
    """查找第k位质数"""
    try:
        if k <= 0:
            return "错误: k必须是正整数"
        
        if k == 1:
            return "2"
            
        # 埃拉托斯特尼筛法
        limit = max(100, k * 20)  # 初始估计足够大的范围
        sieve = [True] * (limit + 1)
        sieve[0] = sieve[1] = False
        
        primes = []
        for num in range(2, limit + 1):
            if sieve[num]:
                primes.append(num)
                if len(primes) == k:
                    return str(num)
                for multiple in range(num * num, limit + 1, num):
                    sieve[multiple] = False
        
        # 如果初始范围不够，扩大范围重新计算
        return nth_prime(k)
    except Exception as e:
        return f"计算错误: {str(e)}"

if __name__ == "__main__":
    mcp.run(transport="stdio")